from langchain.document_loaders import JSONLoader
from faissManager import FAISS
from sent2vec import Sent2VecEmbeddings
from argparse import ArgumentParser

def parse_args():
  parser=ArgumentParser()
  parser.add_argument('--input',type=str,help='指定需要处理的json文件',default='/home/lxy/wxbdata/doc_js.json')
  parser.add_argument('--save_dir',type=str,help='指定索引和key被转成句向量的张量保存路径',default='/data/lxy/RAT/10m21d_onlydoc')
  parser.add_argument('--kv',type=bool,help='用于确定数据库索引的对象是相同键值还是不同键值',default=False,)
  parser.add_argument('--k_name',type=str,help='指定json中哪个键作为key',default='question')
  parser.add_argument('--v_name',type=str,help='指定json中哪个键作为value',default='doc')
  parser.add_argument('--model_name',type=str,help='指定句向量模型,需要是hf格式',default='BAAI/bge-large-zh-v1.5')
  parser.add_argument('--batch_size',type=int,help='编码句向量的batch_size',default=64)
  parser.add_argument('--reuse',type=bool,help='是否重用save_dir保存的张量，不需要重复计算',default=False)
  return parser.parse_args()

if __name__ == '__main__':
  args=parse_args()


  a_documents = JSONLoader(args.input,jq_schema='.[]',content_key=args.v_name).load()
  print('a_documents load success',len(a_documents))

  db = FAISS.from_idxType_and_documents(documents=a_documents,
                                        embedding= Sent2VecEmbeddings(model_name=args.model_name,batch_size=args.batch_size),
                                        tensorSaveDir=args.save_dir,qa=args.kv,reuse=args.reuse)
                                      #   index_type='Flat',do_max_norm=False,normalize_L2=True)


  db.save_local(args.save_dir)